"""
查询构建器

构建ES查询、优化查询性能、分片感知查询等功能。
"""

from datetime import datetime, timezone
from typing import Dict, List, Optional, Any, Union
from .logger import get_module_logger
from .time_utils import TimeRangeProcessor
from .exceptions import ParameterError

logger = get_module_logger(__name__)


class QueryBuilder:
    """ES查询构建器
    
    提供各种ES查询的构建功能，包括时间范围查询、复合查询等。
    """
    
    def __init__(self):
        """初始化查询构建器"""
        self.time_processor = TimeRangeProcessor()
    
    def build_time_range_query(
        self,
        time_field: str,
        start_time: Optional[datetime] = None,
        end_time: Optional[datetime] = None,
        include_lower: bool = True,
        include_upper: bool = True,
        format_str: Optional[str] = None
    ) -> Dict[str, Any]:
        """构建时间范围查询

        Args:
            time_field: 时间字段名
            start_time: 开始时间
            end_time: 结束时间
            include_lower: 是否包含下边界 (默认: True, 使用 gte)
            include_upper: 是否包含上边界 (默认: True, 使用 lte)
            format_str: 时间格式字符串

        Returns:
            时间范围查询
        """
        if not start_time and not end_time:
            return {"match_all": {}}
        
        range_query = {}
        
        if start_time:
            key = "gte" if include_lower else "gt"
            range_query[key] = self._format_time_for_query(start_time, format_str)
        
        if end_time:
            key = "lte" if include_upper else "lt"
            range_query[key] = self._format_time_for_query(end_time, format_str)
        
        if format_str:
            range_query["format"] = format_str
        
        return {
            "range": {
                time_field: range_query
            }
        }
    
    def build_scroll_query(
        self,
        base_query: Dict[str, Any],
        size: int = 1000,
        scroll_timeout: str = "5m",
        sort_fields: Optional[List[Dict[str, Any]]] = None
    ) -> Dict[str, Any]:
        """构建滚动查询
        
        Args:
            base_query: 基础查询
            size: 每批大小
            scroll_timeout: 滚动超时时间
            sort_fields: 排序字段
            
        Returns:
            滚动查询配置
        """
        query = {
            "size": size,
            "query": base_query
        }
        
        if sort_fields:
            query["sort"] = sort_fields
        else:
            # 默认按_doc排序以提高性能
            query["sort"] = ["_doc"]
        
        return query
    
    def build_aggregation_query(
        self,
        base_query: Dict[str, Any],
        aggs: Dict[str, Any],
        size: int = 0
    ) -> Dict[str, Any]:
        """构建聚合查询
        
        Args:
            base_query: 基础查询
            aggs: 聚合配置
            size: 返回文档数量（聚合查询通常设为0）
            
        Returns:
            聚合查询
        """
        return {
            "size": size,
            "query": base_query,
            "aggs": aggs
        }
    
    def build_count_query(self, base_query: Dict[str, Any]) -> Dict[str, Any]:
        """构建计数查询
        
        Args:
            base_query: 基础查询
            
        Returns:
            计数查询
        """
        return {
            "query": base_query
        }
    
    def build_exists_query(self, field: str) -> Dict[str, Any]:
        """构建字段存在查询
        
        Args:
            field: 字段名
            
        Returns:
            存在查询
        """
        return {
            "exists": {
                "field": field
            }
        }
    
    def build_bool_query(
        self,
        must: Optional[List[Dict[str, Any]]] = None,
        should: Optional[List[Dict[str, Any]]] = None,
        must_not: Optional[List[Dict[str, Any]]] = None,
        filter_clauses: Optional[List[Dict[str, Any]]] = None,
        minimum_should_match: Optional[Union[int, str]] = None
    ) -> Dict[str, Any]:
        """构建布尔查询
        
        Args:
            must: 必须匹配的条件
            should: 应该匹配的条件
            must_not: 必须不匹配的条件
            filter_clauses: 过滤条件
            minimum_should_match: 最小should匹配数
            
        Returns:
            布尔查询
        """
        bool_query = {}
        
        if must:
            bool_query["must"] = must
        if should:
            bool_query["should"] = should
        if must_not:
            bool_query["must_not"] = must_not
        if filter_clauses:
            bool_query["filter"] = filter_clauses
        if minimum_should_match is not None:
            bool_query["minimum_should_match"] = minimum_should_match
        
        return {"bool": bool_query}
    
    def build_term_query(self, field: str, value: Any) -> Dict[str, Any]:
        """构建精确匹配查询
        
        Args:
            field: 字段名
            value: 匹配值
            
        Returns:
            精确匹配查询
        """
        return {
            "term": {
                field: value
            }
        }
    
    def build_terms_query(self, field: str, values: List[Any]) -> Dict[str, Any]:
        """构建多值匹配查询
        
        Args:
            field: 字段名
            values: 匹配值列表
            
        Returns:
            多值匹配查询
        """
        return {
            "terms": {
                field: values
            }
        }
    
    def build_wildcard_query(self, field: str, pattern: str) -> Dict[str, Any]:
        """构建通配符查询
        
        Args:
            field: 字段名
            pattern: 通配符模式
            
        Returns:
            通配符查询
        """
        return {
            "wildcard": {
                field: pattern
            }
        }
    
    def combine_queries(self, queries: List[Dict[str, Any]], operator: str = "must") -> Dict[str, Any]:
        """组合多个查询
        
        Args:
            queries: 查询列表
            operator: 组合操作符 (must, should, must_not, filter)
            
        Returns:
            组合查询
        """
        if not queries:
            return {"match_all": {}}
        
        if len(queries) == 1:
            return queries[0]
        
        if operator not in ["must", "should", "must_not", "filter"]:
            raise ParameterError(f"无效的组合操作符: {operator}", "operator", operator)
        
        return {
            "bool": {
                operator: queries
            }
        }
    
    def _format_time_for_query(self, dt: datetime, format_str: Optional[str] = None) -> str:
        """格式化时间用于查询
        
        Args:
            dt: datetime对象
            format_str: 格式字符串
            
        Returns:
            格式化的时间字符串
        """
        if format_str:
            return dt.strftime(format_str)
        else:
            # 默认使用ISO格式
            return dt.isoformat()


class OptimizedQueryBuilder(QueryBuilder):
    """优化的查询构建器
    
    提供查询性能优化功能。
    """
    
    def __init__(self, es_client=None):
        """初始化优化查询构建器
        
        Args:
            es_client: ES客户端（用于获取索引信息）
        """
        super().__init__()
        self.es_client = es_client
    
    def build_optimized_time_range_query(
        self,
        index_pattern: str,
        time_field: str,
        start_time: Optional[datetime] = None,
        end_time: Optional[datetime] = None,
        additional_filters: Optional[List[Dict[str, Any]]] = None
    ) -> Dict[str, Any]:
        """构建优化的时间范围查询
        
        Args:
            index_pattern: 索引模式
            time_field: 时间字段
            start_time: 开始时间
            end_time: 结束时间
            additional_filters: 额外过滤条件
            
        Returns:
            优化的查询
        """
        # 构建基础时间查询
        time_query = self.build_time_range_query(time_field, start_time, end_time)
        
        # 组合额外过滤条件
        filter_clauses = [time_query]
        if additional_filters:
            filter_clauses.extend(additional_filters)
        
        # 使用filter context以提高性能
        optimized_query = {
            "bool": {
                "filter": filter_clauses
            }
        }
        
        return optimized_query
    
    def build_shard_aware_query(
        self,
        base_query: Dict[str, Any],
        routing_field: Optional[str] = None,
        routing_value: Optional[str] = None,
        preference: Optional[str] = None
    ) -> Dict[str, Any]:
        """构建分片感知查询
        
        Args:
            base_query: 基础查询
            routing_field: 路由字段
            routing_value: 路由值
            preference: 查询偏好
            
        Returns:
            分片感知查询配置
        """
        query_config = {
            "query": base_query
        }
        
        # 添加路由信息
        if routing_value:
            query_config["routing"] = routing_value
        
        # 添加查询偏好
        if preference:
            query_config["preference"] = preference
        
        return query_config
    
    def estimate_query_cost(self, query: Dict[str, Any], index: str) -> Dict[str, Any]:
        """估算查询成本
        
        Args:
            query: 查询
            index: 索引名
            
        Returns:
            成本估算
        """
        cost_factors = {
            "complexity_score": self._calculate_query_complexity(query),
            "estimated_docs": "unknown",
            "performance_tips": []
        }
        
        # 如果有ES客户端，获取更详细的信息
        if self.es_client:
            try:
                # 使用封装的count方法估算匹配文档数
                count_result = self.es_client.count(index=index, body={"query": query})
                cost_factors["estimated_docs"] = count_result["count"]

                # 基于文档数量给出性能建议
                if cost_factors["estimated_docs"] > 1000000:
                    cost_factors["performance_tips"].append("考虑添加更多过滤条件以减少匹配文档数")

            except Exception as e:
                logger.warning(f"无法估算查询成本: {e}")
        
        return cost_factors
    
    def _calculate_query_complexity(self, query: Dict[str, Any]) -> int:
        """计算查询复杂度
        
        Args:
            query: 查询
            
        Returns:
            复杂度分数
        """
        complexity = 0
        
        def analyze_clause(clause):
            nonlocal complexity
            
            if isinstance(clause, dict):
                for key, value in clause.items():
                    if key == "bool":
                        complexity += 2
                        if isinstance(value, dict):
                            for bool_key, bool_value in value.items():
                                if isinstance(bool_value, list):
                                    complexity += len(bool_value)
                                    for item in bool_value:
                                        analyze_clause(item)
                    elif key in ["range", "term", "terms", "match", "wildcard"]:
                        complexity += 1
                    elif key in ["nested", "has_child", "has_parent"]:
                        complexity += 5
                    
                    if isinstance(value, (dict, list)):
                        analyze_clause(value)
            elif isinstance(clause, list):
                for item in clause:
                    analyze_clause(item)
        
        analyze_clause(query)
        return complexity
    
    def suggest_query_optimizations(self, query: Dict[str, Any]) -> List[str]:
        """建议查询优化
        
        Args:
            query: 查询
            
        Returns:
            优化建议列表
        """
        suggestions = []
        
        # 检查是否使用了filter context
        if self._has_query_context_only(query):
            suggestions.append("考虑将不需要评分的条件移到filter context中")
        
        # 检查通配符查询
        if self._has_wildcard_queries(query):
            suggestions.append("通配符查询可能影响性能，考虑使用前缀查询或其他替代方案")
        
        # 检查嵌套查询
        if self._has_nested_queries(query):
            suggestions.append("嵌套查询成本较高，确保必要时才使用")
        
        # 检查范围查询
        if self._has_unbounded_range_queries(query):
            suggestions.append("无界范围查询可能影响性能，尽量指定上下界")
        
        return suggestions
    
    def _has_query_context_only(self, query: Dict[str, Any]) -> bool:
        """检查是否只使用了query context"""
        def check_context(clause):
            if isinstance(clause, dict):
                if "bool" in clause:
                    bool_clause = clause["bool"]
                    return "filter" not in bool_clause and ("must" in bool_clause or "should" in bool_clause)
                return True
            return False
        
        return check_context(query)
    
    def _has_wildcard_queries(self, query: Dict[str, Any]) -> bool:
        """检查是否包含通配符查询"""
        def check_wildcard(clause):
            if isinstance(clause, dict):
                if "wildcard" in clause:
                    return True
                for value in clause.values():
                    if isinstance(value, (dict, list)) and check_wildcard(value):
                        return True
            elif isinstance(clause, list):
                return any(check_wildcard(item) for item in clause)
            return False
        
        return check_wildcard(query)
    
    def _has_nested_queries(self, query: Dict[str, Any]) -> bool:
        """检查是否包含嵌套查询"""
        def check_nested(clause):
            if isinstance(clause, dict):
                if "nested" in clause:
                    return True
                for value in clause.values():
                    if isinstance(value, (dict, list)) and check_nested(value):
                        return True
            elif isinstance(clause, list):
                return any(check_nested(item) for item in clause)
            return False
        
        return check_nested(query)
    
    def _has_unbounded_range_queries(self, query: Dict[str, Any]) -> bool:
        """检查是否包含无界范围查询"""
        def check_unbounded_range(clause):
            if isinstance(clause, dict):
                if "range" in clause:
                    range_clause = clause["range"]
                    for field_range in range_clause.values():
                        if isinstance(field_range, dict):
                            has_lower = any(key in field_range for key in ["gte", "gt"])
                            has_upper = any(key in field_range for key in ["lte", "lt"])
                            if not (has_lower and has_upper):
                                return True
                for value in clause.values():
                    if isinstance(value, (dict, list)) and check_unbounded_range(value):
                        return True
            elif isinstance(clause, list):
                return any(check_unbounded_range(item) for item in clause)
            return False
        
        return check_unbounded_range(query)
